import torch
from torch import nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.nn import GATConv, Linear, to_hetero


class Classifier(nn.Module):
    def forward(self, x_A, x_B, index_edge_label):
        print("index_edge_label00000000000000000000000000000000000000000000000000000000",index_edge_label)
        print("x_A.shape",x_A.shape)

        edge_feat_A = x_A[index_edge_label[0]]
        edge_feat_B = x_B[index_edge_label[1]]
        return (edge_feat_A * edge_feat_B).sum(dim=-1)


class GNN(nn.Module):
    def __init__(self, hidden_channels):
        super(GNN, self).__init__()
        self.conv1 = SAGEConv(hidden_channels, hidden_channels)
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x


class Model(nn.Module):
    def __init__(self, hidden_channels):
        super().__init__()
        self.gnn = GNN(hidden_channels)
        self.classifier = Classifier()

    def change_model(self, data):
        self.gnn = to_hetero(self.gnn, metadata=data.metadata())
        # pass

    def forward(self, batch_data):
        """
        aoh_label_index = train_data[("alarm", "on", "host")].edge_label_index
        aoh_edge_label = train_data[("alarm", "on", "host")].edge_label
        atb_label_index = data[("alarm", "to", "bussiness_tree")].edge_label_index
        atb_edge_label = data[("alarm", "to", "bussiness_tree")].edge_label
        hbb_label_index = data[("host", "belongsto", "bussiness_tree")].edge_label_index
        hbb_edge_label = data[("host", "belongsto", "bussiness_tree")].edge_label

        Args:
           111 data:

        Returns:

        """
        batch_data = batch_data.to("cuda:0")
        aoh_label_index = batch_data[("alarm", "on", "host")].edge_label_index
        # aoh_edge_label = batch_data[("alarm", "on", "host")].edge_label
        atb_label_index = batch_data[("alarm", "to", "bussiness_tree")].edge_label_index
        # atb_edge_label = batch_data[("alarm", "to", "bussiness_tree")].edge_label
        hbb_label_index = batch_data[("host", "belongsto", "bussiness_tree")].edge_label_index
        # hbb_edge_label = batch_data[("host", "belongsto", "bussiness_tree")].edge_label
        print("111111111111111111111111111")
        print(batch_data)
        print("2222222222222222222222222222")

        # print(batch_data["alarm"].node_id,"###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data["host"].node_id,"###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data["bussiness_tree"].node_id,"###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data["alarm"].x.shape, "###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data["host"].x.shape, "###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data["bussiness_tree"].x.shape,"###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        #
        # print(aoh_label_index, "1###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(atb_label_index, "2###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(hbb_label_index,"3###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data[("alarm", "rev_on", "host")].edge_label_index, "1###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data[("alarm", "rev_to", "bussiness_tree")].edge_label_index, "2###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
        # print(batch_data[("host", "rev_belongsto", "bussiness_tree")].edge_label_index,"3###################$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")

        x_dict ={
            "alarm": batch_data["alarm"].x,
            "host": batch_data["host"].x,
            "bussiness_tree": batch_data["bussiness_tree"].x
        }
        b_dict = {
            ("alarm", "on", "host"):aoh_label_index,
            ('alarm', 'to', 'bussiness_tree'): atb_label_index,
            ("host", "belongsto", "bussiness_tree"): hbb_label_index
        }
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!",batch_data.edge_index_dict)

        out_label_index_dict = self.gnn(x_dict, batch_data.edge_index_dict)
        # print("out_label_index_dict", out_label_index_dict['alarm'].shape)
        # print("out_label_index_dict", out_label_index_dict['host'].shape)
        # print("out_label_index_dict", out_label_index_dict['bussiness_tree'].shape)

        # aoh_out = self.classifier(out_label_index_dict["alarm"], out_label_index_dict["host"], aoh_label_index)
        atb_out = self.classifier(out_label_index_dict["alarm"], out_label_index_dict["bussiness_tree"],
                                  atb_label_index)
        # hbb_out = self.classifier(out_label_index_dict["host"], out_label_index_dict["bussiness_tree"], hbb_label_index)
        # return aoh_out, atb_out, hbb_out

        return atb_out